package guestbook;

import java.util.*;
import java.util.concurrent.*;

import com.google.appengine.api.datastore.*;
import com.google.appengine.repackaged.com.google.common.collect.Lists;

public class LAsyncDatastoreService implements AsyncDatastoreService
{
	public LAsyncDatastoreService(AsyncDatastoreService bds, Context context)
	{
		this.ds = bds;
		this.context = context;
	}

	private AsyncDatastoreService ds;
	public Context context;

	public void logCall(String method)
	{
		context.tally.tallyCall("async-" + method);
	}

	@Override
	public Collection<Transaction> getActiveTransactions()
	{
		if (Tally.ON)
			logCall("getActiveTransactions");

		return ds.getActiveTransactions();
	}

	@Override
	public Transaction getCurrentTransaction()
	{
		if (Tally.ON)
			logCall("getCurrentTransaction");

		return ds.getCurrentTransaction();
	}

	@Override
	public Transaction getCurrentTransaction(Transaction returnedIfNoTxn)
	{
		if (Tally.ON)
			logCall("getCurrentTransaction");

		return ds.getCurrentTransaction(returnedIfNoTxn);
	}

	@Override
	public PreparedQuery prepare(Query query)
	{
		if (Tally.ON)
			logCall("prepare");

		return ds.prepare(query);
	}

	@Override
	public PreparedQuery prepare(Transaction txn, Query query)
	{
		if (Tally.ON)
			logCall("prepare");

		return ds.prepare(txn, query);
	}

	@Override
	public Future<KeyRange> allocateIds(String kind, long num)
	{
		if (Tally.ON)
			logCall("allocateIds");

		return ds.allocateIds(kind, num);
	}

	@Override
	public Future<KeyRange> allocateIds(Key parent, String kind, long num)
	{
		if (Tally.ON)
			logCall("allocateIds");

		return ds.allocateIds(parent, kind, num);
	}

	@Override
	public Future<Transaction> beginTransaction()
	{
		if (Tally.ON)
			logCall("beginTransaction");

		return ds.beginTransaction();
	}

	@Override
	public Future<Void> delete(Key... keys)
	{
		if (Tally.ON)
			logCall("delete[]");

		return ds.delete(keys);
	}

	@Override
	public Future<Void> delete(Iterable<Key> keys)
	{
		if (Tally.ON)
			logCall("delete[]");

		return ds.delete(keys);
	}

	@Override
	public Future<Void> delete(Transaction txn, Key... keys)
	{
		if (Tally.ON)
			logCall("delete[]");

		return ds.delete(txn, keys);
	}

	@Override
	public Future<Void> delete(Transaction txn, Iterable<Key> keys)
	{
		if (Tally.ON)
			logCall("delete-Iterable");

		return ds.delete(txn, keys);
	}

	@Override
	public Future<Entity> get(Key key)
	{
		if (Tally.ON)
			logCall("get");

		return ds.get(key);
	}

	@Override
	public Future<Map<Key, Entity>> get(Iterable<Key> keys)
	{
		if (Tally.ON)
			logCall("get-Iterable");

		return ds.get(keys);
	}

	@Override
	public Future<Entity> get(Transaction txn, Key key)
	{
		if (Tally.ON)
			logCall("get");

		return ds.get(txn, key);
	}

	@Override
	public Future<Map<Key, Entity>> get(Transaction txn, Iterable<Key> keys)
	{
		if (Tally.ON)
			logCall("get-Iterable");

		return ds.get(txn, keys);
	}

	@Override
	public Future<Key> put(Entity entity)
	{
		if (Tally.ON)
			logCall("put");

		return new FutureWrapper<Key>(ds.put(entity), entity, context);
	}

	@Override
	public Future<List<Key>> put(Iterable<Entity> entities)
	{
		if (Tally.ON)
			logCall("put");

		return new FutureWrapper<List<Key>>(ds.put(entities), entities, context);
	}

	@Override
	public Future<Key> put(Transaction txn, Entity entity)
	{
		if (Tally.ON)
			logCall("put");

		return new FutureWrapper<Key>(ds.put(txn, entity), entity, context);
	}

	@Override
	public Future<List<Key>> put(Transaction txn, Iterable<Entity> entities)
	{
		if (Tally.ON)
			logCall("put");

		return new FutureWrapper<List<Key>>(ds.put(txn, entities), entities, context);
	}

	public static final class FutureWrapper<T> implements Future<T>
	{
		public Future<T> results;
		public Iterable<Entity> entities;
		public Context context;

		public FutureWrapper(Future<T> results, Entity entity, Context context)
		{
			this.results = results;
			this.entities = Lists.newArrayList(entity);
			this.context = context;
		}

		public FutureWrapper(Future<T> results, Iterable<Entity> entities, Context context)
		{
			this.results = results;
			this.entities = entities;
			this.context = context;
		}

		@Override
		public boolean cancel(boolean mayInterruptIfRunning)
		{
			return results.cancel(mayInterruptIfRunning);
		}

		@Override
		public boolean isCancelled()
		{
			return results.isCancelled();
		}

		@Override
		public boolean isDone()
		{
			return results.isDone();
		}

		@Override
		public T get() throws InterruptedException, ExecutionException
		{
			T r = results.get();
			context.putHelper.put(entities);
			return r;
		}

		@Override
		public T get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException
		{
			T r = results.get(timeout, unit);
			context.putHelper.put(entities);
			return r;
		}
	}
}